"""
使用fasttext获取词向量
"""
import fasttext

import config


def build_model(by_word=True):
    """创建模型并保存"""
    # if by_word:
    #     data_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目" \
    #                 r"\案例-chat_service\corpus\dnn\sort/sim_q_cut_by_word.txt"
    # else:
    #     raise NotImplementedError  # 报未实现的异常

    # 创建模型
    # minCount是词频低于这个值的词直接忽略
    # fasttext.train_unsupervised()创建的对象没有predict方法，
    # fasttext.train_supervised()创建的模型有predict方法，加入的数据必须有__label__XX标签，
    # 若用无__label__XX标签的数据进行训练，预测不出结果，结束时会报错
    model = fasttext.train_unsupervised(config.recall_fasttext_model_train_data_path, wordNgrams=1, epoch=10,
                                        minCount=2)
    # 保存模型
    model.save_model(config.recall_fasttext_model_save_path)


# def build_model2(by_word=True):
#     """创建模型并保存"""
#     # if by_word:
#     #     data_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目" \
#     #                 r"\案例-chat_service\corpus\dnn\sort/sim_q_cut_by_word.txt"
#     # else:
#     #     raise NotImplementedError  # 报未实现的异常
#
#     # 创建模型
#     # minCount是词频低于这个值的词直接忽略
#     # fasttext.train_unsupervised()创建的对象没有predict方法，
#     # fasttext.train_supervised()创建的模型有predict方法，加入的数据必须有__label__XX标签
#     model = fasttext.train_supervised(config.recall_fasttext_model_train_data_path, wordNgrams=1, epoch=1,
#                                       minCount=2)
#     # 保存模型
#     return model


def get_model(by_word=True):
    # 加载模型
    if by_word:
        return fasttext.load_model(config.recall_fasttext_model_save_path)
    else:
        raise NotImplementedError


class FasttextVectorizer:
    def __init__(self):
        self.model = get_model(by_word=True)

    def transform(self, sentences):
        """

        :param sentences: [sentence1_str, sentence2_str, ...]
        :return:
        """
        # fasttext.train_unsupervised().get_sentence_vector()传入的是一个字符串，返回的是一个向量
        # 通过fasttext得到的词向量时100维的
        return [self.model.get_sentence_vector(sentence) for sentence in sentences]

    def fit_transform(self, sentences):
        return self.transform(sentences)
